#1D continuous space model for evolution of range edge with and without the evolution of dispersal
# model gives each individual a Gaussian expectation but allow individuals to vary in their gaussian sd
# plan is to compare situations with and without dispersal evolution and the time to breach the gradient

where<-Sys.info()["sysname"]
if (where=="Darwin"){
  setwd("~/evo-dispersal/KBGrad/")
  #system(paste("R CMD SHLIB Density1D.c"))
  system(paste("R CMD SHLIB PointMetrics1D.c"))
}
if (where=="Linux"){
  setwd("/scratch/jc227089/evo-dispersal/KBGrad/")
  #system(paste("R CMD SHLIB Density1D.c"))
  system(paste("R CMD SHLIB -L/usr/lib64/ -L/usr/lib/gcc/x86_64-redhat-linux/4.4.4/ PointMetrics1D.c"))
}
dyn.load("PointMetrics1D.so")




# Initialises a matrix with n individuals
# individuals have a location (spX, spY) and are 
# initilised with gene for p.disp
init.inds<-function(n, spX, Hmean, Dmean, h2H, h2D, VPH, VPD){
	X<-runif(n, -spX, spX) #position
	D<-rnorm(n, Dmean, sqrt(h2D*VPD)) # individual sd
	H<-rnorm(n, Hmean, sqrt(h2H*VPH)) #Habitat match
	DP<-D+rnorm(n, 0, sqrt((1-h2D)*VPD))
	HP<-H+rnorm(n, 0, sqrt((1-h2H)*VPH))
	cbind(X, D, H, DP, HP)		
}

# Defines the optimal phenotype and how this varies through space
# b defines the gradient of the shift in optimum as per K&B
# satable domain defines the distance from the origin before the gradient starts (gives the population a stable core)
hab.space<-function(X, b, stable.domain){
  X[X<stable.domain & X>(-stable.domain)]<-0
  X[X<0]<-X[X<0]+stable.domain
  X[X>0]<-X[X>0]-stable.domain
	X*b
}

# hassel commins population growth where b=1 (contest competition)
# generates expected fecundity
# Might be worth experimenting with a few different growth functions in the end...
hass.comm<-function(lambda, K, N){
	lambda/(1+(lambda-1)*N/K)
}

# Beverton Holt population growth (same as Hass Comins above)
bev.holt<-function(R0, K, N){
  M<-K/(R0-1)
  R0/(1+N/M)
}

# estimates population density at each X as a continuous function
# uses a set bandwidth = 1 (an individual's zone of influence)
# replaced by metrics()
#Nx.prim<-function(X){
#	.Call("dens", R_X=as.double(X), R_n=length(X))
#}

# gets density, mean and (if evovar==TRUE) variance in breeding values
metrics<-function(popmatrix, evovar=TRUE){
	out<-.Call("metrics", R_X=popmatrix[,"X"], R_H=popmatrix[,"H"],
	R_D=popmatrix[,"D"], R_n=nrow(popmatrix), R_ev=as.numeric(evovar))
	colnames(out)<-c("Nx", "MeanD", "MeanH", "SDD", "SDH")
	if (evovar==FALSE) out<-out[,1:3]
	out
}

# gets density, mean and (if evovar==TRUE) variance in breeding values
# parallel version of function
#pmetrics<-function(popmatrix, evovar=TRUE, ncores){
#  out<-.Call("pmetrics", R_X=popmatrix[,"X"], R_H=popmatrix[,"H"],
#             R_D=popmatrix[,"D"], R_n=nrow(popmatrix), R_ev=as.numeric(evovar), 
#             R_ncores=ncores)
#  colnames(out)<-c("Nx", "MeanD", "MeanH", "SDD", "SDH")
#  if (evovar==FALSE) out<-out[,1:3]
#  out
#}


# gets density, mean and (if evovar==TRUE) variance in breeding values
# differs from above in that calculates through space and then approximates back to the individual
ca.metrics<-function(popmatrix, evovar=TRUE){
  bins<-floor(min(popmatrix[,"X"])):ceiling(max(popmatrix[,"X"]))
  cn<-c("Nx", "MeanD", "MeanH", "SDD", "SDH")
  out.sp<-.Call("sum_metrics", R_X=popmatrix[,"X"], R_H=popmatrix[,"H"],
             R_D=popmatrix[,"D"], R_n=nrow(popmatrix), R_bins=bins, R_nbins=length(bins),
             R_ev=as.numeric(evovar), R_bw=1)
  if (evovar==FALSE) {out.sp<-out.sp[,1:3]; cn<-cn[1:3]}
  out.ind<-c()
  for (ii in 1:ncol(out.sp)){
  	out.ind<-cbind(out.ind, approx(x=bins, y=out.sp[,ii], xout=popmatrix[,"X"])$y)
  }
  colnames(out.ind)<-cn
  matrix(out.ind, ncol=length(cn), dimnames=list(NULL, cn))
}

# gets density, mean and (if evovar==TRUE) variance in breeding values
# differs from above in that bw is adjustable and metrics are scored to spatial bins rather than individuals
sum.metrics<-function(popmatrix, evovar=TRUE, bw=10){
  bins<-floor(min(popmatrix[,"X"])):ceiling(max(popmatrix[,"X"]))
  out<-.Call("sum_metrics", R_X=popmatrix[,"X"], R_H=popmatrix[,"H"],
             R_D=popmatrix[,"D"], R_n=nrow(popmatrix), R_bins=bins, R_nbins=length(bins),
             R_ev=as.numeric(evovar), R_bw=bw)
  colnames(out)<-c("Nx", "MeanD", "MeanH", "SDD", "SDH")
  if (evovar==FALSE) out<-out[,1:3]
  out<-cbind(X=bins, out)
  out
}

#determines a survival probability based on distance moved
#d.surv<-function(dists, k=0.2){
#  exp(-k*dists)
#}

# gives lambda as a function of individual H and locality Optimum.
# Taken from Kirk and Barton 1997
#might be better off with initial formulation as this one becomes negative!!
#h.fit<-function(lambda, H, Opt, Vs){
#  lambda-(H-Opt)^2/(2*Vs)
#}

# Alternative function defining the strength of stabilising selection
h.surv<-function(n=1, H, Opt, k=2){
  n*exp(-k*abs(H-Opt))
}

VE<-function(h, VP){
	(1-h)*VP
}

# function to calculate genetic means and variances weighted for distance
# replaced by metrics()
#gvar.dist<-function(popmatrix, bw=1){
#  w<-outer(X=popmatrix[,"X"], Y=popmatrix[,"X"], FUN="-") #pairwise dists
#  w<-dnorm(x=w, mean=0, sd=bw) #weighted by bw
#  w<-sweep(w, MARGIN=1, STATS=apply(w, 1, sum), FUN="/") #and normalised
#  Hmean<-as.vector(w%*%popmatrix[,"H"]) #vector of mean H
#  names(Hmean)<-rep("A", length(Hmean))
#  Dmean<-as.vector(w%*%popmatrix[,"D"]) #vector of mean D
#  Hvar<-apply(outer(X=Hmean, Y=popmatrix[,"H"], FUN="-")^2*w, 1, sum) #weighted sum of squared deviations
#  Dvar<-apply(outer(X=Dmean, Y=popmatrix[,"D"], FUN="-")^2*w, 1, sum)
#  list(Hmean, Dmean, Hvar, Dvar)
#}

# reproduces individuals based on density, kills parents
# kills offspring based on habitat
# kills dispersing offspring based on dispersal cost
# disperses individuals based on inherited P
repro.disp<-function(popmatrix, K, lambda, b, h2H, h2D, VPH, VPD, evovar=TRUE, stable.domain=spX){
	SDVED<-sqrt(VE(h2D, VPD)) # define environmental variances
	SDVEH<-sqrt(VE(h2H, VPH))
	# calculate traits through space
	mets<-ca.metrics(popmatrix, evovar)
	hab<-hab.space(X=popmatrix[,"X"], b, stable.domain) # habitat returned to each individual
	
	# reproduction
	hab<-h.surv(1, popmatrix[,"HP"], hab)
	offs<-rpois(nrow(popmatrix), hab*hass.comm(lambda, K, mets[,"Nx"])) # surviving offspring
	#offs[is.nan(offs)]<-0 #catch moments when lambda==0, which R can't cope with.
	if (sum(offs)==0) return(NULL) #catch extinctions
	inds<-1:length(offs) #vector indexes for offspring
	inds<-rep(inds, times=offs)
	popmatrix<-popmatrix[inds,] #replace parents with offspring. Offspring inherit location and P from parents
	mets<-matrix(mets[inds,], ncol=ncol(mets), dimnames=list(NULL, colnames(mets)))
	if (class(popmatrix)=="numeric") popmatrix<-matrix(popmatrix, ncol=5, dimnames=list(NULL, c("X", "D", "H", "DP", "HP"))) #cases where only one individual left
	
	# calculate new G and P
	if (evovar==FALSE){
		SDD<-sqrt(h2D*VPD)
		SDH<-sqrt(h2H*VPH)
	} else {
		SDD<-mets[,"SDD"]
		SDH<-mets[,"SDH"]
	}
	
	popmatrix[,"D"]<-rnorm(nrow(popmatrix), mets[,"MeanD"], SDD)
	popmatrix[,"DP"]<-popmatrix[,"D"]+rnorm(nrow(popmatrix), 0, SDVED)
	popmatrix[,"H"]<-rnorm(nrow(popmatrix), mets[,"MeanH"], SDH)
	popmatrix[,"HP"]<-popmatrix[,"H"]+rnorm(nrow(popmatrix), 0, SDVEH)
	
	# Disperse offspring
	disp<-rnorm(nrow(popmatrix), mean=popmatrix[,"X"], sd=exp(popmatrix[,"D"])) #disperse offspring
  	#surv<-rbinom(n=length(disp), size=1, prob=d.surv(abs(disp-popmatrix[,"X"])))
  	popmatrix[,"X"]<-disp
	popmatrix
}

# calculates each individual's fitness
w<-function(popmatrix, K, lambda, b, stable.domain){
	mets<-ca.metrics(popmatrix, evovar=FALSE)
	hab<-hab.space(X=popmatrix[,"X"], b, stable.domain) # habitat returned to each individual
	hab<-h.surv(1, popmatrix[,"HP"], hab)
	hab*hass.comm(lambda, K, mets[,"Nx"])
}


# plots population density and mean trait values through space
plotter.mean<-function(popmatrix, K, lambda, b, stable.domain, filename=NULL, ...){
	popmatrix<-popmatrix[order(popmatrix[,"X"]),]
	summ<-sum.metrics(popmatrix, TRUE, ...)
	hab<-hab.space(X=summ[,"X"], b, stable.domain) # habitat returned to each location
	hab<-h.surv(1, summ[,"MeanH"], hab)
  	if (!is.null(filename)) pdf(file=filename, width=5, height=10)
	par(mfrow=c(3, 2))
  	plot(summ[,"X"], summ[,"MeanH"], xlab="Location", ylab="Mean trait value", pch=19, col="darkorange", type="l")
  	lines(summ[,"X"], hab.space(X=summ[,"X"], b, stable.domain))	
  	plot(summ[,"X"], summ[,"SDH"], xlab="Location", ylab="SD of trait value", pch=19, col="darkorange", type="l")
	
	plot(summ[,"X"], summ[,"MeanD"], xlab="Location", ylab="Mean dispersal value", pch=19, col="darkgreen", type="l")
	plot(summ[,"X"], summ[,"SDD"], xlab="Location", ylab="SD of Dispersal", pch=19, col="darkgreen", type="l")
	
	plot(summ[,"X"], summ[,"Nx"], xlab="Location", ylab="Population density", type="l")
	
	plot(summ[,"X"], summ[,"Nx"]*exp(summ[,"MeanD"])*hab*hass.comm(lambda, K, summ[,"Nx"]), xlab="Location", ylab="Migrant output", type="l")
	if (!is.null(filename)) dev.off()
}


#calculates local dz dx and averages over cline space
dzdx3<-function(popmatrix, K, lambda, b, stable.domain, generation){
  mets<-sum.metrics(popmatrix, evovar=FALSE, bw=1) #calculated to spatial bins
  wf<-function(K, lambda, b, stable.domain){ #fitness
    hab<-hab.space(X=mets[,"X"], b, stable.domain) # habitat returned to each location
    hab<-h.surv(1, mets[,"MeanH"], hab)
    hab*hass.comm(lambda, K, mets[,"Nx"])
  }
  zx<-function(z.vec, x.vec){
    ord<-order(x.vec)
    dx<-c(x.vec[ord][-1], NA)-x.vec[ord]
    dz<-c(z.vec[ord][-1], NA)-z.vec[ord]
    sset<-(is.finite(dx) & is.finite(dz)) #catches instances of infinite values from log of zero density
    mean(dz[sset]/dx[sset], na.rm=TRUE)
  }
  
  left<-which(mets[,"X"]<(-stable.domain))
  right<-which(mets[,"X"]>(stable.domain))
  side<-list(left, right)
  w.tot<-wf(K, lambda, b, stable.domain)
  corr<-c(-1,1)
  out<-c()
  for(ss in 1:2){
    if (length(side[[ss]])<=1) next
    dNx<-corr[ss]*zx(mets[side[[ss]],"Nx"], mets[side[[ss]],"X"])
    dH<-corr[ss]*zx(mets[side[[ss]],"MeanH"], mets[side[[ss]],"X"])
    dD<-corr[ss]*zx(mets[side[[ss]],"MeanD"], mets[side[[ss]],"X"])  
    wdx<-corr[ss]*zx(log(w.tot[side[[ss]]]), mets[side[[ss]],"X"])
    asymm<-corr[ss]*zx(mets[side[[ss]],"Nx"]*exp(mets[side[[ss]],"MeanD"])*w.tot[side[[ss]]], mets[side[[ss]],"X"])
    asymm.l<-corr[ss]*zx(log(mets[side[[ss]],"Nx"]*exp(mets[side[[ss]],"MeanD"])*w.tot[side[[ss]]]), mets[side[[ss]],"X"])
    out<-rbind(out, cbind(dNx, dH, dD, wdx, asymm, asymm.l, generation, ss))
  }
  if (is.null(out)) return(NULL)
  colnames(out)<-c("dNx", "dH", "dD", "wdx", "asymm", "asymm.l", "generation", "ss")
  row.names(out)<-NULL
  out
}

	
# iterates pop over ngens and collects resulting dvec
# then re-initialises population, using dvec, and runs with and without dispersal evolution
mother<-function(n, spX, K, lambda, b, ngens, Hmean, Dmean, h2H, h2D, VPH, VPD, evovar, stable.domain=spX){
  pop<-init.inds(n=n, spX=spX, Hmean=Hmean, Dmean=Dmean,
	h2H=h2H, h2D=h2D, VPH=VPH, VPD=VPD)
	ztrace<-c()
  for(ii in 1:ngens){
    pop<-repro.disp(popmatrix=pop, K=K, lambda=lambda, b=b, 
    	h2H=h2H, h2D=h2D, VPH=VPH, VPD=VPD, evovar=evovar, stable.domain=stable.domain)
    if (is.null(pop)) break
    if (ii>1000){
    	ztrace<-rbind(ztrace, dzdx3(pop, K, lambda, b, stable.domain, ii))
    }
  }
  out<-list(parameters=list(ngens=ii, K=K, lambda=lambda, b=b, n=n, spX=spX,
            ngens=ngens, Hmean=Hmean, Dmean=Dmean, h2H=h2H, h2D=h2D, VPH=VPH,
            VPD=VPD, evovar=evovar), pop=pop, ztrace=ztrace)
  out
}

# iterates pop over ngens and collects resulting dvec
# space is wrapped so there is no gradient and no edges!
mother.wrapped<-function(n, spX, K, lambda, b, ngens, Hmean, Dmean, h2H, h2D, VPH, VPD, evovar, stable.domain=spX){
  pop<-init.inds(n=n, spX=spX, Hmean=Hmean, Dmean=Dmean,
                 h2H=h2H, h2D=h2D, VPH=VPH, VPD=VPD)
  for(ii in 1:ngens){
    pop<-repro.disp(popmatrix=pop, K=K, lambda=lambda, b=b, 
                    h2H=h2H, h2D=h2D, VPH=VPH, VPD=VPD, evovar=evovar, stable.domain=stable.domain)
    pop[pop[,"X"]>spX, "X"]<-pop[pop[,"X"]>spX, "X"]-2*spX
    pop[pop[,"X"]<(-spX), "X"]<-pop[pop[,"X"]<(-spX), "X"]+2*spX
    if (is.null(pop)) break
  }
  out<-list(parameters=list(ngens=ii, K=K, lambda=lambda, b=b, n=n, spX=spX,
                            ngens=ngens, Hmean=Hmean, Dmean=Dmean, h2H=h2H, h2D=h2D, VPH=VPH,
                            VPD=VPD, evovar=evovar), pop=pop)
  out
}